# @Time : 2020/12/4 9:34
# @Author : Fioman 
# @Phone : 13149920693
"""
各种算子测试,差分算子
"""
import os

import imutils
import cv2 as cv
import numpy as np
from 公司算法研究.operator_train.my_logger import log_init


def image_show(name, image, time, showState=False):
    if showState:
        h, w = image.shape[:2]
        rate = 6
        widthShow = int(w / rate)
        heightShow = int(h / rate)
        showImage = cv.resize(image, (widthShow, heightShow), cv.INTER_AREA)
        cv.namedWindow("{}".format(name), cv.WINDOW_AUTOSIZE)
        cv.imshow("{}".format(name), showImage)
        cv.waitKey(time)


def get_width_and_height(fileName):
    """
    根据文件名获取图片的长宽
    :param fileName:
    :return:返回文件的长宽,单位像素
    """
    if fileName[-4:] != ".bmp":
        return None, None
    if "w_" not in fileName or "h_" not in fileName:
        return None, None
    wIndex = fileName.index("w_")
    hIndex = fileName.index("h_")
    width = fileName[wIndex + 2:wIndex + 6]
    height = fileName[hIndex + 2:hIndex + 6]
    return int(width), int(height)


def get_center_by_contour(contour):
    """
    根据轮廓获取中心点坐标,中心点坐标必须是整型的
    :param contour:
    :return:
    """
    (cx, cy), _, _ = cv.minAreaRect(contour)
    return int(cx), int(cy)


def get_rect_points_clockwise(rectPoints):
    """
    根据传过来的四个点,对它们按照顺时针的方向进行排序
    :param rectPoints:
    :return:
    """
    rectPointsSorted = sorted(rectPoints, key=np.sum)
    leftTop = rectPointsSorted[0]
    rightBottom = rectPointsSorted[-1]

    if rectPointsSorted[1][0] > rectPointsSorted[2][0]:
        rightTop = rectPointsSorted[1]
        leftBottom = rectPointsSorted[2]
    else:
        rightTop = rectPointsSorted[2]
        leftBottom = rectPointsSorted[1]

    return leftTop, rightTop, rightBottom, leftBottom


def find_total_board(image, width, height):
    """
    根据宽高还有原始图像获取物料区域,并且将物料区域提取出来.这里想一下,要不要旋转,如果不旋转的话,直接提取有什么问题.
    :param image:
    :param width:
    :param height:
    :return:物料区域,旋转后的物料区域,物料的宽和高
    """
    thresUsed = np.mean(image[:150, :1100]) + 15
    filterRate = 0.2
    T, board = cv.threshold(image, thresUsed, 255, cv.THRESH_BINARY)
    cnts = cv.findContours(board.copy(), cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)

    # 1> 然后将找到的所有满足要求的区域,通过白线连接起来,形成整个物料框
    singleBoardArea = width * height
    largeContours = [cnt for cnt in cnts if cv.contourArea(cnt) > singleBoardArea * filterRate]

    # 2> 如果满足条件的区域大于2,就要连线,然后重新再找一次
    if len(largeContours) >= 2:
        centerPoints = [get_center_by_contour(contour) for contour in largeContours]
        for index, center in enumerate(centerPoints):
            anotherCenter = centerPoints[(index + 1) % len(centerPoints)]
            cv.line(board, center, anotherCenter, 255, 20)
        image_show("JointCenterLine", board, 0, True)

        # 3> 连线之后再找一遍
        cnts = cv.findContours(board.copy(), cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
        cnts = imutils.grab_contours(cnts)
        cntFirstFinded = sorted(cnts, key=cv.contourArea, reverse=True)[0]
    else:
        cntFirstFinded = largeContours[0]

    # 这个轮廓找到以后,然后将这个轮廓的四个角的影响去除掉
    rectPoints = np.int0(cv.boxPoints(cv.minAreaRect(cntFirstFinded)))
    leftTop, rightTop, rightBottom, leftBottom = get_rect_points_clockwise(rectPoints)
    boardFinal = image.copy()
    moveWidth = int(width * 0.3)
    moveHeight = int(height * 0.3)
    extend = int(8 * mm2Pix)
    boardFinal[leftTop[1] - extend:leftTop[1] + moveHeight, leftTop[0] - extend:leftTop[0] + moveWidth] = 0
    boardFinal[rightTop[1] - extend:rightTop[1] + moveHeight, rightTop[0] - moveWidth:rightTop[0] + extend] = 0
    boardFinal[leftBottom[1] - extend - moveHeight:leftBottom[1], leftBottom[0] - extend:leftBottom[0] + moveWidth] = 0
    boardFinal[rightBottom[1] - moveHeight - extend:rightBottom[1], rightBottom[0] - moveWidth:rightBottom[0] + extend] = 0
    image_show("FinalBoardWithoutAngle", boardFinal, 0, True)

    # 最后找到轮廓
    cnts = cv.findContours(board.copy(), cv.RETR_LIST, cv.CHAIN_APPROX_SIMPLE)
    cnts = imutils.grab_contours(cnts)
    cntFinalFinded = sorted(cnts, key=cv.contourArea, reverse=True)[0]

    box = cv.minAreaRect(cntFinalFinded)
    (cx, cy), (w, h), angle = box
    if angle > 70:
        angle = 90 - 70
        box = (cx, cy), (h, w), angle
    if angle < -70:
        angle = 90 + angle
        box = (cx, cy), (h, w), angle

    rectPoints = np.int0(cv.boxPoints(box))
    boxMask = np.zeros(image.shape, dtype=np.uint8)
    # 将轮廓画出来
    imageColor = cv.cvtColor(image, cv.COLOR_GRAY2BGR)
    cv.drawContours(boxMask, [rectPoints], -1, 255, -1)
    imageGray = cv.bitwise_and(image, boxMask)  # 将物料框以外的地方黑化
    cv.drawContours(imageColor, [rectPoints], -1, (0, 255, 0), 3)
    image_show("TotalBoardContour", imageColor, 0, True)

    # 然后就是截取整个物料的区域.
    widthHalf = int(w / 2)
    heightHalf = int(h / 2)
    yStart = max(int(cy - heightHalf), 0)  # 之所以有零的情况,是因为旋转之后物料框可能跑外面去了
    yEnd = min(int(cy + heightHalf), image.shape[0])  # y也是一样的,旋转之后图跑外面去了,整个时候结束位置应该是图像的高度
    xStart = max(int(cx - widthHalf), 0)
    xEnd = min(int(cx + widthHalf), image.shape[1])
    M = cv.getRotationMatrix2D((cx, cy), angle, 1)
    totalBoard = cv.warpAffine(image, M, (image.shape[1], image.shape[0]))[yStart:yEnd, xStart:xEnd]
    jointLineMask = np.zeros(image.shape, dtype=np.uint8)
    jointPos = 1283
    jointLineMask[:, jointPos - 20:jointPos + 20] = 255
    jointLineMaskRotated = cv.warpAffine(jointLineMask, M, (image.shape[1], image.shape[0]))[yStart:yEnd, xStart:xEnd]
    image_show("TotalBoard", totalBoard, 0, True)

    return totalBoard, w, h, angle, jointLineMaskRotated


def get_two_boxes(image, width, height, bigThres, smallThres, mm2Pix):
    if image is None:
        return "error", "图像异常,请重拍"
    if len(image.shape) == 3:
        image = cv.cvtColor(image, cv.COLOR_BGR2GRAY)

    # 1.先获取整个物料区域.
    totalBoard, boardW, boardH, angle, jointLine = find_total_board(image, width, height)
    totalBoardColor = cv.cvtColor(totalBoard, cv.COLOR_GRAY2BGR)

    # 想将图片分成两个图片
    upDownSlice = int(40 * mm2Pix)
    upBoard = totalBoard.copy()[:int(boardH / 2 + upDownSlice), :]
    downBoard = totalBoard.copy()[:int(boardH / 2 - upDownSlice), :]
    upJointLine = jointLine.copy()[:int(boardH / 2 + upDownSlice), :]
    downJointLine = jointLine.copy()[:int(boardH / 2 - upDownSlice), :]

    # 分别去求这个两个板的内板信息
    upRes, upInfo, upBox = get_inner_board_info(upBoard, width, height, bigThres, smallThres, "up", mm2Pix)
    downRes, downInfo, downBox = get_inner_board_info(downBoard, width, height, bigThres, smallThres, "down", mm2Pix)


def find_edge_by_DOG(board):
    """
    使用差分获取边缘信息,高斯差分算法,利用两个高斯模糊差分,看看边缘的信息.
    :param board:
    :return:
    """
    cv.destroyAllWindows()
    gaussainR = 5
    gaussainBlurred1 = cv.GaussianBlur(board, (0, 0), gaussainR)
    gaussainBlurred2 = cv.GaussianBlur(board, (0, 0), gaussainR * 1.6)
    gaussainDif = cv.subtract(gaussainBlurred1, gaussainBlurred2)
    image_show("GaussainDif", gaussainDif, 0, True)
    return gaussainDif


def find_edge_by_sobel(board):
    sobelX = cv.Sobel(board,-1,1,0,ksize=3)
    sobelY = cv.Sobel(board,-1,1,0,ksize=3)
    sobelXY = cv.addWeighted(sobelX,0.5,sobelY,0.5,0)
    image_show("sobelX",sobelX,0,True)
    image_show("sobelY",sobelY,0,True)
    image_show("sobelXY",sobelXY,0,True)


def find_edge_by_scharr(board):
    scharrX = cv.Scharr(board,-1,1,0)
    scharrY = cv.Scharr(board,-1,0,1)
    scharrXY = cv.addWeighted(scharrX,0.5,scharrY,0.5,0)
    image_show("ScharrX",scharrX,0,True)
    image_show("scharrY",scharrY,0,True)
    image_show("scharrXY",scharrXY,0,True)


def find_edge_by_roberts(board):
    kernelX = np.array([[-1,0],[0,1]],dtype=int)
    kernelY = np.array([[0,-1],[1,0]],dtype=int)
    robertX = cv.filter2D(board,cv.CV_16S,kernelX)
    robertY = cv.filter2D(board,cv.CV_16S,kernelY)

    robertX = cv.convertScaleAbs(robertX)
    robertY = cv.convertScaleAbs(robertY)
    robertXY = cv.addWeighted(robertX,0.5,robertY,0.5,0)
    image_show("RobertX",robertX,0,True)
    image_show("RobertY",robertY,0,True)
    image_show("RobertXY",robertXY,0,True)


def find_edge_by_prewitt(board):
    pass


def find_edge_by_laplacian(board):
    blurred = cv.GaussianBlur(board,(0,0),1)
    lapGray = cv.Laplacian(blurred,cv.CV_16S,ksize=3)
    lapGray = cv.convertScaleAbs(lapGray)
    image_show("LapGray",lapGray,0,True)
    addLap = cv.add(blurred,lapGray)

    scharrX = cv.Scharr(addLap,cv.CV_16S,1,0)
    scharrX = cv.convertScaleAbs(scharrX)
    scharrY = cv.Scharr(addLap,cv.CV_16S,0,1)
    scharrY = cv.convertScaleAbs(scharrY)
    scharrXY = cv.addWeighted(scharrX,0.5,scharrY,0.5,0)
    image_show("ScharrXY",scharrXY,0,True)


def get_inner_board_info(board, width, height, bigThres, smallThres, boardPos, mm2Pix):
    """
    根据尺寸获取内板的信息,测试各种算子分步骤进行.
    :param board:
    :param width:
    :param height:
    :param bigThres:
    :param smallThres:
    :param boardPos:
    :param mm2Pix:
    :return:
    """
    moveWidth = max(int(width * 0.38), int(150 * mm2Pix))
    moveHeight = int(200 * mm2Pix)
    boardColor = cv.cvtColor(board, cv.COLOR_GRAY2BGR)

    # 使用方差去获取边缘信息.
    # edgeImg = find_edge_by_DOG(board)
    # edgeImg = find_edge_by_sobel(board)
    # edgeImg = find_edge_by_scharr(board)
    # edgeImg = find_edge_by_roberts(board)
    # edgeImg = find_edge_by_prewitt(board)
    edgeImg = find_edge_by_laplacian(board)



if __name__ == '__main__':
    calcThresBig = 15
    calcThresSmall = 10
    mm2Pix = 3.25

    testFilePath = r"D:\operator_raw"
    fileNames = []
    logger = log_init(__name__)
    if not os.path.isdir(testFilePath):
        fileName, testFilePath = os.path.split(testFilePath)

    resOkPath = os.path.join(testFilePath, "ok")
    resFailPath = os.path.join(testFilePath, "fail")
    if not os.path.exists(resOkPath):
        os.mkdir(resOkPath)
    if not os.path.exists(resFailPath):
        os.mkdir(resFailPath)

    for root, dirs, files in os.walk(testFilePath):
        if root != testFilePath:  # 过滤掉不是要测试的主目录
            continue
        for file in files:
            fileNames.append(file)

    for index, fileName in enumerate(fileNames):
        filePathReal = os.path.join(testFilePath, fileName)
        widthTest, heightTest = get_width_and_height(fileName)
        if widthTest is None or heightTest is None:
            logger.debug("*" * 10 + "第{} 张图 ({}),图像名称错误,已跳过".format(index + 1, fileName) + "*" * 10)
            continue
        logger.debug("*" * 10 + " 第 {} 张图 ({})".format(index + 1, fileName) + "*" * 10)

        imgTest = cv.imread(filePathReal, cv.IMREAD_GRAYSCALE)
        totalBoardTest = None
        res, info, upBox, downBox, totalBox = get_two_boxes(imgTest, widthTest, heightTest, calcThresBig, calcThresSmall, mm2Pix)
